import os
import torch
import cv2
import numpy as np
from torch.utils.data import Dataset

class VideoDataset(Dataset):
    def __init__(self, dataset_path,images_path,clip_len):
        self.dataset_path=dataset_path#数据集的地址
        self.split=images_path#训练集测试集验证集的名字
        self.clip_len=clip_len#数据的深度就是照片的张数，有的有50多张。后续有操作从中选取16张

        #后续数据预处理的值
        self.resize_height=128
        self.resize_width=171
        self.crop_size=112

        #读取对应训练集、验证集、测试集下的各种类别的行为动作
        #每个行为动作下的视频已经被处理成单个图片数据
        #将对应动作的数据的文件名作为标签保存到labels列表中，对应的动作数据集的路径保存到self.fnames列表中，标签和数据是一一对应的状态
        folder=os.path.join(self.dataset_path,images_path)#images_path是train/val/test
        '''这种写法中，self.fnames 被定义为类实例的属性，而 labels 被定义为局部变量。
        self.fnames 可以在类的任何方法中访问，因为它属于类实例。
        labels 只能在当前方法的作用域内访问，不能在类的其他方法中使用。'''
        self.fnames,labels=[],[]#fname是每个视频帧图片的文件夹名字。label是这个文件夹的上一级，就是ApplyEyeMakeup，
        for label in sorted(os.listdir(folder)):#folder就是train/val、test。label就是标签
            #fname 变量会依次取 os.listdir() 返回的列表中的每一个元素。这个列表包含的是指定目录下的所有文件和子目录的名字
            for fname in os.listdir(os.path.join(folder,label)):
                self.fnames.append(os.path.join(folder,label,fname))#fname现在是train文件夹下v_ApplyEyeMakeup_g01_c01的路径，这个子文件夹下就是图片
                labels.append(label)
        print('Number of {} videos: {:d}'.format(images_path,len(self.fnames)))

        #获取对应视频的标签，并将标签转化为int的数字类型，同时转化为array格式
        self.label2index={label:index for index,label in enumerate(sorted(set(labels)))}#label是标签，index for...是0,1,2...
        self.label_array=np.array([self.label2index[label] for label in labels],dtype=int)#每个标签都被替换为其在 self.label2index 字典中的索引值

    def __len__(self):
        return len(self.fnames)
    def __getitem__(self, index):
        buffer=self.load_frames(self.fnames[index])#加载对应类别的动作数据集，并转化为（帧数，高，宽，3通道）
        buffer = self.crop(buffer,self.clip_len,self.crop_size)#在数据集的深度高度宽度进行随机裁剪，将数据转化为(clip_len,112,112,3)
        buffer=self.normalize(buffer)#对模型惊醒归一化处理
        buffer=self.to_tensor(buffer)#对维度进行转化

        #获取对应视频的标签数据
        labels=np.array(self.label_array[index])

        #返回torch格式的特征和标签
        '''torch.from_numpy 函数用于将一个 NumPy 数组转换为 PyTorch 张量（Tensor）。
        这个函数创建的张量与原 NumPy 数组共享内存，这意味着对其中一个数组的数据进行修改会反映到另一个数组中'''
        return torch.from_numpy(buffer),torch.from_numpy(labels)


    def  load_frames(self,file_dir):#传的参数是fnames，就是每个视频帧图片的文件夹。是一个列表，存的地址
        #将文件夹下的数据集进行排序
        frames=sorted([os.path.join(file_dir,img) for img in os.listdir(file_dir)])
        #获取该动作数据的长度
        frame_count=len(frames)#其中一个是41
        #生成一个空的（frame_count,resize_height,resize_width,3)维度的数据
        buffer =np.empty((frame_count,self.resize_height,self.resize_width,3),np.dtype('float'))#全部都是0，维度是（41,128,171,3）
        for i,frame_name in enumerate(frames):
            #利用cv去读取图片数据，并转化为np.array格式
            frame = np.array(cv2.imread(frame_name)).astype(np.float64)#第一章图片的像素矩阵
            #不断遍历循环赋值给buffer
            buffer[i]=frame
        return buffer#（41,128,171,3）

    def crop(self,buffer,clip_len,crop_size):#clip_len=16,crop_size=112
        time_index=np.random.randint(buffer.shape[0]-clip_len)#生成一个深度方向上的随机长度.buffer第一个维度第一个文件夹是41，随机生成0-25的整形
        height_index=np.random.randint(buffer.shape[1]-crop_size)#生成一个高度方向上的随机长度的起始值
        width_index=np.random.randint(buffer.shape[2]-crop_size)#生成一个宽度方向上的随机长度的起始值

        #利用切片在视频上进行提取，获得一个（clip_len,112,112,3）的数据
        buffer =buffer[time_index:time_index+clip_len,
                height_index:height_index+crop_size,
                width_index:width_index+crop_size,:]
        return buffer

    def normalize(self,buffer):#归一化函数，影响不大
        #进行归一化
        for i,frame in enumerate(buffer):
            frame -=np.array([[90.0,98.0,102.0]])#可能是均值
            buffer[i]=frame
        return buffer

    def to_tensor(self,buffer):#tensor函数
        #进行维度的转化，将最后的一个维度调转到第一维
        return buffer.transpose((3,0,1,2))

if __name__ == '__main__':
    from torch.utils.data import DataLoader
    train_data=VideoDataset(dataset_path='data/ucf101',images_path='train',clip_len=16)#clip_len=16是取16张
    train_dataloader = DataLoader(train_data,batch_size=64,shuffle=True,num_workers=0)

    val_data = VideoDataset(dataset_path='data/ucf101', images_path='val', clip_len=16)  # clip_len=16是取16张
    val_dataloader = DataLoader(train_data, batch_size=64, shuffle=True, num_workers=0)

    train_data = VideoDataset(dataset_path='data/ucf101', images_path='test', clip_len=16)  # clip_len=16是取16张
    train_dataloader = DataLoader(train_data, batch_size=64, shuffle=False, num_workers=0)


